// SPDX-License-Identifier: GPL-2.0 OR Linux-OpenIB
/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: Support reg umm mr
 */

#include <linux/mm.h>
#include <linux/dma-mapping.h>
#include <linux/sched/signal.h>
#include <linux/sched/mm.h>
#include <linux/export.h>
#include <linux/hugetlb.h>
#include <linux/slab.h>
#include <linux/jiffies.h>
#include <linux/kmemleak.h>
#include <rdma/ib_umem.h>
#include <trace/hooks/rdma_pangea.h>

#include "uverbs.h"

static int __ib_umm_init_sg_table(struct ib_umem *umem,
	phys_addr_t phys_addr, unsigned int sgl_count, size_t align_size)
{
	int ret;
	struct scatterlist *sg;
	phys_addr_t curr_addr;
	size_t pg_off;
	unsigned int i, len;

	ret = sg_alloc_table(&umem->sg_head, sgl_count, GFP_KERNEL);
	if (ret) {
		ofed_log_print(PANGEA_ERR, "Alloc scatterlist table failed, ret: %d.", ret);
		return ret;
	}

	sg = umem->sg_head.sgl;
	pg_off = phys_addr & (_UL(SZ_4M) - 1);
	curr_addr = phys_addr;

	for (i = 0; i < sgl_count; i++) {
		len = SZ_2G - pg_off;
		if (i == (sgl_count - 1))
			len = phys_addr + align_size - curr_addr;

		sg_dma_address(sg) = curr_addr;
		sg_dma_len(sg) = len;
		curr_addr += len;
		sg = sg_next(sg);
		pg_off = 0;
	}

	return 0;
}

static unsigned int __ib_umm_get_sgl_counts(phys_addr_t phys_addr, size_t align_size)
{
	size_t pg_off = phys_addr & (_UL(SZ_4M) - 1);

	return DIV_ROUND_UP(align_size + pg_off, SZ_2G);
}

struct ib_umem *ib_umm_get(struct ib_device *device, unsigned long addr,
	size_t size, int access)
{
	struct ib_umem *umem;
	unsigned long npages;
	int ret;
	unsigned int sgl_count;
	size_t align_size = 0;
	phys_addr_t phys_addr;
	struct timespec64 start;

	trace_ofed_monitor_exec_time_start(&start);

	if (!size) {
		ofed_log_print(PANGEA_ERR, "Bad MR size.");
		return ERR_PTR(-EINVAL);
	}

	phys_addr = (phys_addr_t)ib_umm_get_pyhs_addr(addr);
	if (phys_addr == (phys_addr_t)addr) {
		ofed_log_print(PANGEA_ERR, "Get phys address from umm failed, addr(0x%lx).", addr);
		return ERR_PTR(-EINVAL);
	}

	umem = kzalloc(sizeof(*umem), GFP_KERNEL);
	if (!umem) {
		ofed_log_print(PANGEA_ERR, "No memory, alloc size %zu.", sizeof(*umem));
		return ERR_PTR(-ENOMEM);
	}

	umem->is_umm_mem  = 1;
	umem->ibdev      = device;
	umem->length     = size;
	umem->address    = addr;
	/*
	 * Drivers should call ib_umem_find_best_pgsz() to set the iova
	 * correctly.
	 */
	umem->iova = addr;
	umem->writable   = ib_access_writable(access);
	umem->owning_mm = current->mm;

	npages = ib_umem_num_pages(umem);
	align_size = npages * PAGE_SIZE;
	sgl_count = __ib_umm_get_sgl_counts(phys_addr & PAGE_MASK, align_size);

	ret = __ib_umm_init_sg_table(umem, phys_addr & PAGE_MASK, sgl_count, align_size);
	if (ret)
		goto out;

	umem->sg_nents = umem->sg_head.nents;
	umem->nmap = sgl_count;
	ofed_log_print(PANGEA_INFO,
		"Reg umm mr, addr(0x%lx), phys addr(0x%llx), sgl_count(%u), npages(%lu), length(%zu).",
		addr, (unsigned long long)phys_addr, sgl_count, npages, size);
out:

	if (ret < 0)
		kfree(umem);

	trace_ofed_monitor_exec_time_end(&start, device, "Get umm info");

	return ret < 0 ? ERR_PTR(ret) : umem;
}
EXPORT_SYMBOL(ib_umm_get);

void ib_umm_release(struct ib_umem *umem)
{
	sg_free_table(&umem->sg_head);
	kfree(umem);
}
EXPORT_SYMBOL(ib_umm_release);
